#!/bin/bash

voxceleb1_path=/ceph/home/zhangy20/datasets/CN-Celeb1
voxceleb2_path=~/datasets/CN-Celeb2
musan_path=~/datasets/musan
rirs_path=~/datasets/RIRS_NOISES

vox1_trials_path=data/cnc1_trials.npy

embedding_save_dir=npy

. ./path.sh

stage=2
echo stage $stage

# format data dir structure by soft link
if [ $stage -eq 0 ];then
    if [ ! -d data/wav_files ]; then
        mkdir -p data/wav_files
    fi

    # link and format dir structure
    rm -rf data/wav_files/dev
    mkdir -p data/wav_files/dev
    for spk in `cat ${CN_celeb1_path}/dev/dev.lst`;do
        ln -s ${CN_celeb1_path}/data/${spk} data/wav_files/dev/cnc1-$spk
    done

    for spk in `ls ${CN_celeb2_path}/data`;do
        ln -s ${CN_celeb2_path}/data/${spk} data/wav_files/dev/cnc2-$spk
    done

    rm -rf data/wav_files/eval
    mkdir -p data/wav_files/eval
    for type in enroll test;do
        ln -s $CN_celeb1_path/eval/$type data/wav_files/eval/$type
    done

    # format read CN-Celeb1 trials and save to binary npy
    rm -rf data/cnc1_trials.npy
    python $OPENASV_ROOT/scripts/format_cnc1_trials.py \
        --src_trials_path $CN_celeb1_path/eval/lists/trials.lst \
        --dst_trials_path $cnc1_trials_path \
        --npy_dir $eval_spk_embedding_save_path
fi


# build data list
if [ $stage -eq 1 ];then
    echo build dev data list
    python3 $OPENASV_ROOT/scripts/build_datalist.py \
        --extension wav \
        --dataset_dir data/wav_files/dev \
        --data_list_path data/dev_list

    python3 $OPENASV_ROOT/scripts/split_train_val.py \
        --input_data_list_path data/dev_list \
        --output_train_list_path data/train_list \
        --output_val_list_path data/val_list

    echo build eval data list
    python3 $OPENASV_ROOT/scripts/build_datalist.py \
        --extension wav \
        --dataset_dir data/wav_files/eval \
        --data_list_path data/eval_list

    echo build musan data list
    python3 $OPENASV_ROOT/scripts/build_datalist.py \
        --extension wav \
        --dataset_dir $musan_path \
        --data_list_path data/musan_list

    echo build rirs data list
    python3 $OPENASV_ROOT/scripts/build_datalist.py \
        --extension wav \
        --dataset_dir $rirs_path \
        --data_list_path data/rirs_list
fi


if [ $stage -eq 2 ];then
    CUDA_VISIBLE_DEVICES=0,1,2,3 python3 $OPENASV_ROOT/main.py \
        --batch_size 800 \
        --num_workers 40 \
        --save_top_k 5 \
        --train_list_path data/train_list \
        --val_list_path data/val_list \
        --test_list_path data/eval_list \
        --musan_list_path data/musan_list \
        --rir_list_path data/rirs_list \
        --embedding_save_dir $embedding_save_dir \
        --max_frames 400 \
        --learning_rate 0.001 \
        --gpus 4 \
		--distributed_backend dp
fi


if [ $stage -eq 3 ];then
    rm -rf npy/enroll
    mkdir -p npy/enroll
    rm -rf npy/test
    mkdir -p npy/test

    ckpt_path=ckpt.pt
    CUDA_VISIBLE_DEVICES=0 python3 $OPENASV_ROOT/main.py \
        --batch_size 1 \
        --batch_size 200 \
        --num_workers 20 \
        --save_top_k 5 \
        --train_list_path data/train_list \
        --val_list_path data/val_list \
        --test_list_path data/eval_list \
        --musan_list_path data/musan_list \
        --rir_list_path data/rirs_list \
        --max_frames 200 \
        --learning_rate 0.001 \
        --gpus 1 \
        --embedding_save_dir $embedding_save_dir \
        --resume_from_checkpoint $ckpt_path \
        --test

    for file in `ls npy/enroll`;do
        mv npy/enroll/$file npy/enroll/${file%-*}-enroll.npy
    done
fi

if [ $stage -eq 4 ];then
    python $OPENASV_ROOT/score.py \
        --trials_path $cnc1_trials_path
fi
